# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/5 15:14 
@Author  : ZhangShenao 
@File    : 1.加载文档并存储至Pinecone向量数据库.py 
@Desc    : 加载文档并存储至Pinecone向量数据库
"""
import dotenv
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载环境变量
dotenv.load_dotenv()

# 使用文档加载器加载文档
loader = UnstructuredMarkdownLoader(file_path="./docs/项目API资料.md", mode="single")
docs = loader.load()

# 使用RecursiveCharacterTextSplitter,对文档进行分割
separators = [
    '\n\n',
    '\n',
    r'。|！|？',
    r'\.\s|\!\s|\?\s',  # 英文标点符号后面通常需要加空格
    r'；|;\s',
    r'，|,\s',
    ' ',
    ''
]
splitter = RecursiveCharacterTextSplitter(
    separators=separators,  # 指定自定义分隔符列表
    chunk_size=100,  # 指定文档分块大小
    chunk_overlap=20,  # 指定文档块之间的重叠部分大小
    add_start_index=True,  # 添加文档块的起始索引
    is_separator_regex=True,  # 指定分隔符为正则表达式形式
)
chunks = splitter.split_documents(docs)

# 创建Pinecone向量数据库
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)

# 将切割后的文档保存到Pinecone
vector_store.add_documents(chunks)

print("保存文档成功")
